import torch
import numpy as np
from torch.utils.data import DataLoader,Sampler
from lib.datasets.kitti.kitti_dataset import KITTI_Dataset
import math
from lib.datasets.nuscenes import NuScenesDataset

class PaddingSampler(Sampler):
    def __init__(self, data_source, batch_size: int):
        num=len(data_source)
        self.num=num
        self.idx_list=list(range(num))
        self.bz=batch_size
        self.idx_list+=list(range(math.ceil(num/batch_size)*batch_size-num))

    def __iter__(self):
        return iter(self.idx_list)

    def __len__(self):
        return int(math.ceil(self.num/self.bz))


# init datasets and dataloaders
def my_worker_init_fn(worker_id):
    np.random.seed(np.random.get_state()[1][0] + worker_id)


def build_dataloader(cfg, workers=4):
    # perpare dataset
    cfg_test=cfg.copy()
    cfg_test.update({
        'type':'KITTI',
        'root_dir':'data/KITTIDataset',
        'resolution':[1280,384],
        'writelist':['Car'],
    })
    test_set = KITTI_Dataset(split=cfg['test_split'], cfg=cfg_test)
    train_set = NuScenesDataset(split=cfg['train_split'], cfg=cfg)

    # prepare dataloader
    test_sampler=PaddingSampler(test_set,cfg['batch_size'])
    train_loader = DataLoader(dataset=train_set,
                              batch_size=cfg['batch_size'],
                              num_workers=workers,
                              worker_init_fn=my_worker_init_fn,
                              shuffle=True,
                              pin_memory=False,
                              drop_last=True)
    test_loader = DataLoader(dataset=test_set,
                             batch_size=cfg['batch_size'],
                             num_workers=workers,
                             worker_init_fn=my_worker_init_fn,
                             shuffle=False,
                             pin_memory=False,
                             drop_last=False,sampler=test_sampler)

    return train_loader, test_loader
